import os
import torch
import torch.utils.data
import torchvision
import numpy as np
import pandas as pd
from data.apple_dataset import AppleDataset
import utility.utils as utils
import utility.transforms as T
import datetime
import time
import matplotlib.pyplot as plt
import torch
import torch.utils.data
from PIL import Image
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from data.apple_dataset import AppleDataset
from utility.engine import train_one_epoch, evaluate

import utility.utils as utils
import utility.transforms as T

def get_transform(train):
    transforms = []
    # 修改为仅处理图像的转换逻辑
    transforms.append(lambda image: T.ToTensor()(image, target=None)[0])  # 只取图像部分
    if train:
        transforms.append(lambda image: T.RandomHorizontalFlip(0.5)(image, target=None)[0])
    
    def transform_fn(image):
        for t in transforms:
            image = t(image)
        return image
    
    return transform_fn


def save_binary_mask(mask, output_path, img_name):
    """
    将二值化掩码张量保存为黑白图片
    参数:
        mask: numpy数组或torch张量 (H, W) 或 (1, H, W)
        output_path: 输出目录
        img_name: 原始图片文件名
    """
    os.makedirs(output_path, exist_ok=True)
    
    # 转换为numpy数组并去除多余维度
    if isinstance(mask, torch.Tensor):
        mask = mask.squeeze().cpu().numpy()
    else:
        mask = mask.squeeze()
    
    # 二值化处理 (0-255)
    binary_mask = (mask > 0.5).astype(np.uint8) * 255
    
    # 创建PIL图像并保存
    mask_img = Image.fromarray(binary_mask)
    save_path = os.path.join(output_path,"binary",img_name)

    mask_img.save(save_path)
    print(f"Saved binary mask to {save_path}")


def save_mask_as_image(mask, output_path, img_name, mask_id, data_path):
    if mask.dtype != np.float32:
        mask = mask.astype(np.float32)
    if mask.min() < 0 or mask.max() > 1:
        mask = (mask - mask.min()) / (mask.max() - mask.min())
    print(mask.shape)
    if len(mask.shape) == 4:
        mask = mask[0, 0, :, :]
    elif len(mask.shape) == 3:
        mask = mask[0, :, :]

    mask = torch.from_numpy(mask)
    os.makedirs(output_path, exist_ok=True)

    mask_path = os.path.join(output_path, f"{img_name}_mask_{mask_id}.png")
    try:

        original_img_path = os.path.join(data_path, f"{img_name}")  # 假设原始图像是jpg格式
        compare_path = os.path.join(output_path, "compared")
        os.makedirs(compare_path, exist_ok=True)
        print(original_img_path)
        if os.path.exists(original_img_path):
            # 读取原图
            original_img = plt.imread(original_img_path)

            # 获取原图尺寸
            height, width = original_img.shape[0], original_img.shape[1]
            
            # 创建动态尺寸的画布（根据原图尺寸计算英寸）
            fig, ax = plt.subplots(figsize=(width/100, height/100), dpi=100)  # 100 dpi时1像素=0.01英寸
            plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

            # 显示原图（使用实际尺寸）
            ax.imshow(original_img, extent=[0, width, 0, height])  # 动态设置显示范围
            ax.axis('off')
            
            # 显示掩码叠加效果（匹配原图尺寸）
            ax.imshow(mask.numpy().squeeze(), 
                     cmap='Reds', 
                     alpha=0.3, 
                     extent=[0, width, 0, height])  # 使用实际宽高

            # 保存图片（保持1:1像素比例）
            compare_img_path = os.path.join(compare_path, img_name)
            plt.savefig(compare_img_path, 
                       pad_inches=0, 
                       dpi=100, 
                       bbox_inches='tight')
            plt.close()
            print(f"Saved comparison to {compare_img_path}")
    except Exception as e:
        print(f"Error saving mask: {e}")


# 在save_mask_as_image函数之后添加新函数
def draw_boxes_on_image(original_path, boxes, output_path, img_name):
    try:
        # 创建输出目录
        os.makedirs(output_path, exist_ok=True)
        
        # 打开原始图片
        img = Image.open(original_path).convert("RGB")

        plt.figure(figsize=(img.width/100, img.height/100), dpi=100)
        plt.imshow(img)
        
        # 绘制每个预测框（使用转换后的坐标）
        for box in boxes:
            x1, y1, x2, y2 = box
            rect = plt.Rectangle((x1, y1), x2 - x1, y2 - y1,  # 修正参数顺序
                               linewidth=2, edgecolor='lime', facecolor='none')
            plt.gca().add_patch(rect)
        
        # 保存结果
        plt.axis('off')
        output_file = os.path.join(output_path,"withBox", img_name)
        plt.savefig(output_file, bbox_inches='tight', pad_inches=0)
        plt.close()
        
        print(f"Saved box image to {output_file}")
    except Exception as e:
        print(f"Error drawing boxes: {e}")


def predict_with_gui(data_path, weight_file, model_type='mrcnn', device='cuda', single_image=None):
    """
    GUI版本的预测函数
    参数:
        data_path: 测试数据路径
        weight_file: 模型权重文件路径
        model_type: 模型类型 ('mrcnn' 或 'frcnn')
        device: 计算设备 ('cuda' 或 'cpu')
        single_image: 要处理的单张图片文件名
    返回:
        results: 包含预测结果的列表，每个元素是一个字典，包含图像名称和预测框信息
    """

    # 初始化设备
    device = torch.device(device if torch.cuda.is_available() else 'cpu')

    # 加载模型
    print("Loading model...")
    model = torch.load(weight_file, map_location=device)
    model.to(device)
    model.eval()

    # 设置 mask 输出路径
    output_mask_path = "static/SaveLoad"
    os.makedirs(output_mask_path, exist_ok=True)

    # 直接加载单张图片
    img_path = os.path.join(data_path, single_image)
    if not os.path.exists(img_path):
        raise FileNotFoundError(f"图片文件 {img_path} 不存在")

    # 加载并预处理图片
    img = Image.open(img_path).convert("RGB")
    transform = get_transform(train=False)
    img_tensor = transform(img).unsqueeze(0).to(device)

    # 开始预测
    print("Starting prediction...")
    with torch.no_grad():
        # 模型预测
        outputs = model(img_tensor)

        # 处理预测结果
        output = outputs[0]
        boxes = output['boxes'].detach().cpu().numpy()
        scores = output['scores'].detach().cpu().numpy()
        masks = output['masks'].detach().cpu().numpy()

        # 置信度过滤 (阈值设为0.5)
        keep_idx = scores > 0.5
        boxes = boxes[keep_idx]
        scores = scores[keep_idx]
        masks = masks[keep_idx]

        # 合并所有mask
        combined_mask = None
        for mask in masks:
            binary_mask = (mask[0] > 0.5).astype(np.uint8) * 255
            if combined_mask is None:
                combined_mask = binary_mask
            else:
                combined_mask = np.maximum(combined_mask, binary_mask)

        # 保存合并结果
        # 在预测结果处理部分添加：
        if combined_mask is not None:
        # 保存叠加效果图
            save_mask_as_image(combined_mask, output_mask_path, single_image, "combined", data_path)  #//combine 完全没有用但是我不想删
            # 新增：保存黑白掩码图
            save_binary_mask(combined_mask, output_mask_path, single_image)
        if len(boxes) > 0:
            # 绘制带框的原图
            draw_boxes_on_image(
                original_path=img_path,
                # mask=combined_mask,
                boxes=boxes.copy(),
                output_path=os.path.join(output_mask_path),
                img_name=single_image
            )
        # 返回结果
        result = {
            'image_name': single_image,
            'boxes': boxes,
            'scores': scores
        }

    print("Prediction completed!")
    return [result]


def main():
#     # 固定参数设置
    args = {
        'data_path': r"./",
        'output_file': "./output/1.txt",
        'weight_file': "./myapple101.pth",
        'device': "cuda",
        'model_type': 'mrcnn'  # 直接指定模型类型
    }
    print("Initializing prediction...")
#     # 执行预测
    results = predict_with_gui(
        data_path=args['data_path'],
        weight_file=args['weight_file'],
        model_type=args['model_type'],
        device=args['device'],

        single_image = "234.png"
    )
#
#     # 保存结果到文件
#     with open(args['output_file'], 'w') as f:
#         for result in results:
#             boxes = result['boxes'].astype(int)
#             scores = result['scores']
#             for box, score in zip(boxes, scores):
#                 line = f"{result['image_name']},{box[0]},{box[1]},{box[2]},{box[3]},{score:.4f}\n"
#                 f.write(line)
#     print(f"预测结果已保存至 {args['output_file']}")
#     df_list = []
#     for result in results:
#         for box, score in zip(result['boxes'], result['scores']):
#             df_list.append([
#                 result['image_name'],
#                 *box.astype(int).tolist(),
#                 round(float(score), 4)
#             ])
#
#     pred_df = pd.DataFrame(df_list,
#                            columns=['image_name', 'x1', 'y1', 'x2', 'y2', 'confidence'])
#
#     # 保存结果
#     pred_df.to_csv(args['output_file'], index=False)
#     print(f"Saved predictions to {args['output_file']}")
#
if __name__ == "__main__":
    main()
